import os
from PIL import Image
import torch
from torch.utils.data import DataLoader
from torchvision import models, transforms  # transforms图形预处理的模块

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))

"""图像的预处理
 - 将输入图像的大小缩放为256x256的尺寸
 - 在图像的中心位置进行裁剪，将图像裁剪为224x224的尺寸
 - 将图像转换为PyTorch中的Tensor对象
 - 对图像进行标准化操作，使得图像像素的数值范围符合神经网络的输入要求
"""
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


# 载入模型
def load_model():
    """
    第二行创建了一个ResNet-34模型对象，并将该模型的预训练权重加载到模型中。
    第三行将该模型对象移动到指定的计算设备上（GPU或CPU）进行计算。
    第四行将该模型对象设置为评估模式（即在推理时对输入数据进行前向传递）。
    """
    model = models.resnet34(pretrained=True)
    model.to(device)
    model.eval()  # 用来提取特征，所以使用评估模式
    return model


# 定义 特征提取器
def feature_extract(model, x):
    """
    特征提取器
    对输入数据张量进行第一层卷积操作，并将结果保存在x中。
    对x进行批标准化操作，并将结果保存在x中。
    对x进行ReLU激活操作，并将结果保存在x中。
    对x进行最大池化操作，并将结果保存在x中。
    分别对x进行模型的layer1、layer2、layer3和layer4四个卷积层的操作。
    对x进行平均池化操作，并将结果保存在x中。
    将x展平成一个一维向量，并返回结果。

    :param model:
    :param x:
    :return:
    """
    x = model.conv1(x)
    x = model.bn1(x)
    x = model.relu(x)
    x = model.maxpool(x)

    x = model.layer1(x)
    x = model.layer2(x)
    x = model.layer3(x)
    x = model.layer4(x)

    x = model.avgpool(x)  # 平均池化下采样
    x = torch.flatten(x, 1)  # 展平得到一维特征向量

    return x


